package com.wkh.datastructure.AVL;
//AVL树
public class AVLTree {

    static class AVLNode{
        public int key;
        public Object value;
        public AVLNode left;
        public AVLNode right;
        int height = 1;//初始化值为1，刚创建的树的节点都为1

        public AVLNode(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        public AVLNode(int key) {
            this.key = key;
        }

        public AVLNode(int key, Object value, AVLNode left, AVLNode right) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
        }
    }

    //获取高度
    private int height(AVLNode node){
        return node == null ? 0 : node.height;
    }
    //更新节点的高度(新增、删除、旋转)
    private void updateHeight(AVLNode node){
        node.height = Integer.max(height(node.left),height(node.right));
    }

    /**
     *
     * @return 0 1 -1 是平衡的
     *         >1 <-1 是不平衡的
     */
    //平衡因子(balance factor) = 左子树高度 - 右子树高度
    private int bf(AVLNode node){
        return height(node.left) - height(node.right);
    }

    /**
     *
     * @param node 不平衡的节点
     * @return 返回新的节点
     */
    //右旋
    private AVLNode rightRotate(AVLNode node){
        AVLNode left = node.left;
        AVLNode right = left.right;
        left.right = node; //上位
        node.left = right; //换爹
        updateHeight(left);
        updateHeight(right);
        return left;
    }
    /**
     *
     * @param node 不平衡的节点
     * @return 返回新的节点
     */
    //左旋
    private AVLNode leftRotate(AVLNode node){
        AVLNode right = node.right;
        AVLNode left = right.left;
        right.left = node;
        node.right = left;
        updateHeight(right);
        updateHeight(left);
        return right;
    }

    /**
     * 先左旋不平衡节点的左子树，再右旋不平衡的节点
     * @param node 不平衡的节点
     * @return 返回新的节点
     */
    private AVLNode leftRightRotate(AVLNode node){
        //左旋
        leftRotate(node.left);
        //右旋
        return rightRotate(node);
    }
    /**
     * 先右旋不平衡节点的右子树，再左旋不平衡的节点
     * @param node 不平衡的节点
     * @return 返回新的节点
     */
    private AVLNode rightLeftRotate(AVLNode node){
        //右旋
        rightRotate(node.right);
        //左旋
        return leftRotate(node);
    }
    //检查节点是否失衡，重新平衡代码
    private AVLNode balance(AVLNode node){
        if (node == null){
            return null;
        }
        int bf = bf(node);
        if (bf>1 && bf(node.left)>=0){//RR
            return rightRotate(node);
        }else if (bf>1 && bf(node.left) < 0){//RL
            return rightLeftRotate(node);
        }else if (bf < -1 && bf(node.right)>0){//LR
            return leftRightRotate(node);
        }else if (bf<-1 && bf(node.right) <= 0){//LL
            return leftRotate(node);
        }
        return node;
    }
    AVLNode root;
    //新增或添加节点
    public void put(int key,Object value){
         root = doPut(root,key,value);
    }
    //新增或添加节点
    private AVLNode doPut(AVLNode node,int key,Object value){
        //1.找到空位，创建新结点
        if (node == null){
            return new AVLNode(key,value);
        }
        //2.key存在，更新节点的值
        if (key == node.key){
            node.value = value;
        }
        //3.继续查找
        if (key<node.key){
            node.left = doPut(node.left,key,value);//向左找
        }else {
            node.right = doPut(node.right,key,value);//向右找
        }
        //更新节点高度
        updateHeight(node);
        //重新检查平衡
        return balance(node);
    }
    //删除节点
    public void remove(int key){
        doRemove(root,key);
    }

    private AVLNode doRemove(AVLNode node, int key) {
        //1.node == null
        if (node == null){
            return null;
        }
        //2.没有找到key
        if (key < node.key){
            node.left =  doRemove(node.left,key);
        }else if (key > node.key){
            node.right = doRemove(node.right,key);
        }else{
            //3.找到key  ①没有子节点
            //          ②只有一个子节点
            //          ③有两个子节点
            if (node.left == null && node.right == null){
                return null;
            }else if (node.left == null){
                node = node.right;
            }else if (node.right == null){
                node = node.left;
            }else{
                //找到该节点的后继节点
                //s是后继节点
                AVLNode s = node.right;
                while (s.left != null){
                    s = s.left;
                }
                s.right = doRemove(node.right, s.key);
                s.left = node.left;
                node = s;
            }
        }
        //4.更新高度
        updateHeight(node);
        //5.重新检查是否失衡
        return balance(node);
    }
}
